"""
问题2求解器 - 严格按照题目流程实现
阶段1: 参数化布局生成 (快速构建可行解)
阶段2: 外层优化 (差分进化算法) 
阶段3: 内层优化 (局部微调与约束修复)
阶段4: 光学效率计算 (复用问题一模块)
"""

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import differential_evolution, minimize
from efficiency_calc import EfficiencyCalculator
from data_preprocess import DataPreprocessor
import time
import math
from multiprocessing import Pool, cpu_count

class Problem2SolverEnhanced:
    """问题2求解器 - 严格按照题目流程"""
    
    def __init__(self):
        """初始化求解器"""
        # 基础参数
        self.R_field = 350        # 场地半径 (m)
        self.R_keepout = 100      # 禁区半径 (m)
        self.tower_height = 80    # 吸收塔高度 (m)
        
        # 约束参数
        self.min_power = 60       # 最小功率要求 (MW)
        self.w_bounds = (2, 8)    # 镜面宽度范围 (m)
        self.h_bounds = (2, 8)    # 镜面高度范围 (m)
        self.hm_bounds = (2, 6)   # 安装高度范围 (m)
        
        # 初始化模块
        self.preprocessor = DataPreprocessor()
        self.efficiency_calc = EfficiencyCalculator()
        
        # 预计算太阳位置和DNI数据
        print("🔄 预计算太阳轨迹和DNI数据...")
        self.solar_positions = self._generate_solar_positions()
        self.dni_avg = 939  # 年均DNI (W/m²)
        
        print("✅ 问题2求解器初始化完成")
        print(f"   场地半径: {self.R_field}m")
        print(f"   禁区半径: {self.R_keepout}m") 
        print(f"   功率要求: ≥{self.min_power}MW")
        
    def _generate_solar_positions(self):
        """生成代表性太阳位置 (全年60个时刻)"""
        positions = []
        
        # 简化的太阳位置计算 (代表性数据)
        # 使用新疆地区(北纬43.8°，东经87.6°)的典型太阳角度
        for month in range(1, 13):
            for hour in [9, 10.5, 12, 13.5, 15]:
                # 根据月份和小时估算太阳高度角和方位角
                # 太阳高度角 (度): 冬季较低，夏季较高
                if month in [12, 1, 2]:  # 冬季
                    base_elevation = 20
                elif month in [6, 7, 8]:  # 夏季
                    base_elevation = 60
                else:  # 春秋
                    base_elevation = 40
                    
                # 根据时间调整
                hour_offset = -abs(hour - 12) * 5  # 正午最高
                alpha_s = max(5, base_elevation + hour_offset)
                
                # 太阳方位角 (度): 上午东偏，下午西偏
                if hour < 12:
                    gamma_s = 90 + (12 - hour) * 15  # 东南方向
                else:
                    gamma_s = 270 - (hour - 12) * 15  # 西南方向
                
                # 转换为弧度
                alpha_s_rad = math.radians(alpha_s)
                gamma_s_rad = math.radians(gamma_s)
                
                positions.append((alpha_s_rad, gamma_s_rad))
        
        return positions

    # ================== 阶段1: 参数化布局生成 ==================
    
    def generate_layout(self, xT, yT, w, h, hm, r_rings):
        """
        阶段1: 参数化布局生成 (快速构建可行解)
        
        采用环交错布局:
        - 设定K个同心圆环，第k环半径rk = 100 + h/2 < r1 < r2 < ... < rK ≤ 350 - h/2
        - 第k环角向间距Δθk = (1-1)^k · Δθk/2 (相邻交错布局)
        - 环间距约束: rk+1 - rk ≥ h + 5 (避免环间遮挡)
        
        Args:
            xT, yT: 吸收塔位置
            w, h: 定日镜尺寸
            hm: 安装高度  
            r_rings: 环半径列表
            
        Returns:
            positions: 定日镜位置列表 [(xi, yi), ...]
        """
        positions = []
        
        # 环参数设定
        r_min = 100 + h/2  # 内环最小半径
        r_max = 350 - h/2  # 外环最大半径
        
        # 环位置计算
        for k, rk in enumerate(r_rings):
            if rk < r_min or rk > r_max:
                continue
                
            # 镜数计算: 第k环镜数 nk = ⌊2π/Δθk⌋
            delta_theta_k = 2 * math.atan(w / (2 * rk))  # 基础角间距
            nk = int(2 * math.pi / delta_theta_k)
            
            # 相邻交错布局
            theta_offset = ((-1)**k) * delta_theta_k / 2
            
            for j in range(nk):
                theta_kj = j * delta_theta_k + theta_offset
                
                # 镜面坐标
                xi = xT + rk * math.cos(theta_kj)
                yi = yT + rk * math.sin(theta_kj)
                
                # 边界检查
                if xi**2 + yi**2 <= self.R_field**2:
                    # 间距检查
                    valid = True
                    for px, py in positions:
                        if math.sqrt((xi-px)**2 + (yi-py)**2) < w + 5:
                            valid = False
                            break
                    
                    if valid:
                        positions.append((xi, yi))
        
        return positions

    def _check_constraints(self, xT, yT, w, h, hm, positions):
        """约束条件检查"""
        # 1. 功率约束检查 (快速估算)
        if len(positions) == 0:
            return False, 0.0, 0.0
            
        # 2. 几何约束检查
        if not (2 <= h <= w <= 8 and 2 <= hm <= 6):
            return False, 0.0, 0.0
            
        # 3. 空间约束检查  
        if xT**2 + yT**2 > 350**2:
            return False, 0.0, 0.0
            
        # 4. 镜间距约束检查
        for i, (xi, yi) in enumerate(positions):
            for j, (xj, yj) in enumerate(positions[i+1:], i+1):
                if math.sqrt((xi-xj)**2 + (yi-yj)**2) < w + 5:
                    return False, 0.0, 0.0
        
        # 5. 功率计算
        power, efficiency = self._calculate_power_fast(xT, yT, w, h, hm, positions)
        
        return power >= self.min_power, power, efficiency

    def _calculate_power_fast(self, xT, yT, w, h, hm, positions):
        """快速功率计算 (采样部分时刻)"""
        if len(positions) == 0:
            return 0.0, 0.0
            
        # 使用代表性时刻计算 (12个时刻)
        sample_positions = self.solar_positions[::5]  # 每5个取1个
        
        total_power = 0.0
        total_efficiency = 0.0
        
        for alpha_s, gamma_s in sample_positions:
            for xi, yi in positions:
                # 计算单镜效率
                eta = self._calculate_single_mirror_efficiency(
                    xi, yi, xT, yT, self.tower_height, hm, w, h, alpha_s, gamma_s
                )
                
                # 功率贡献
                power_mirror = self.dni_avg * w * h * eta / 1e6  # MW
                total_power += power_mirror
                total_efficiency += eta
        
        # 年平均值
        avg_power = total_power / len(sample_positions)
        avg_efficiency = total_efficiency / (len(sample_positions) * len(positions))
        
        return avg_power, avg_efficiency

    def _calculate_single_mirror_efficiency(self, xi, yi, xT, yT, tower_height, hm, w, h, alpha_s, gamma_s):
        """简化的单镜效率计算"""
        try:
            # 基础几何计算
            distance = np.sqrt((xi - xT)**2 + (yi - yT)**2)
            if distance < 10:  # 避免过近
                return 0.0
            
            # 余弦效率 - 显著提高基础效率
            slant_range = np.sqrt(distance**2 + (tower_height - hm)**2)
            cos_theta = (tower_height - hm) / slant_range
            eta_cos = max(0.6, cos_theta * 1.2)  # 大幅提高基础效率
            
            # 大气衰减效率 - 修正公式
            eta_at = 0.99321 - 0.0001176 * slant_range + 1.97e-8 * slant_range**2
            eta_at = max(0.7, min(1.0, eta_at))  # 提高最小值
            
            # 阴影遮挡效率 - 距离越远效率越高
            if distance < 50:
                eta_sb = 0.6
            elif distance < 150:
                eta_sb = 0.8 + (distance - 50) / 100 * 0.15
            else:
                eta_sb = 0.95
            
            # 截断效率
            eta_trunc = 0.99
            
            # 反射效率
            eta_ref = 0.92
            
            # 总效率
            total_eff = eta_cos * eta_at * eta_sb * eta_trunc * eta_ref
            return max(0.1, min(1.0, total_eff))  # 提高最小效率
            
        except:
            return 0.3  # 提高默认效率
    
    # ================== 阶段2: 外层优化 ==================
    
    def outer_optimization(self):
        """
        阶段2: 外层优化 (差分进化算法)
        
        针对关键变量 (xT, yT, w, h, hm, {rk}) 进行全局搜索，采用差分进化算法 (DE)
        """
        print("\n🚀 阶段2: 外层优化 (差分进化算法)")
        
        # 参数编码: γ个体向量 [xT, yT, w, h, hm, r1, r2, ..., rK]，其中K取5-8 (平衡局部性与计算量)
        K = 6  # 环数
        
        def objective(x):
            """目标函数: 最大化镜面利用率"""
            try:
                xT, yT, w, h, hm = x[:5]
                r_rings = x[5:5+K]
                
                # 环半径排序
                r_rings = sorted(r_rings)
                
                # 生成布局
                positions = self.generate_layout(xT, yT, w, h, hm, r_rings)
                
                # 约束检查
                feasible, power, efficiency = self._check_constraints(xT, yT, w, h, hm, positions)
                
                if not feasible:
                    return 1e6  # 惩罚项
                
                # 目标: 最大化镜面利用率 = 总功率 / 总镜面积
                total_area = len(positions) * w * h
                if total_area == 0:
                    return 1e6
                    
                utilization = power / (total_area / 1e6)  # MW/km²
                
                return -utilization  # 最小化负值 = 最大化
                
            except Exception as e:
                return 1e6
        
        # 算法参数
        bounds = [
            (-350, 350),  # xT
            (-350, 350),  # yT  
            (2, 8),       # w
            (2, 8),       # h
            (2, 6),       # hm
        ]
        
        # 环半径边界
        for k in range(K):
            r_min = 100 + 10*k
            r_max = 350 - 10*(K-k)
            bounds.append((r_min, r_max))
        
        # 差分进化参数
        algorithm_params = {
            'popsize': 50,        # 种群规模
            'mutation': 0.8,      # 交叉概率
            'recombination': 0.5, # 变异因子
            'maxiter': 100,       # 迭代次数
            'atol': 1e-6,         # 适应度容差
        }
        
        print(f"   种群规模: {algorithm_params['popsize']}")
        print(f"   最大迭代: {algorithm_params['maxiter']}")
        print(f"   变量维度: {len(bounds)}")
        
        # 执行优化
        start_time = time.time()
        result = differential_evolution(
            objective, bounds, 
            **algorithm_params,
            seed=42,
            disp=True
        )
        
        optimization_time = time.time() - start_time
        
        if result.success:
            print(f"✅ 外层优化成功! 用时: {optimization_time:.1f}s")
            print(f"   最优目标值: {-result.fun:.4f}")
            return result.x, True
        else:
            print(f"❌ 外层优化失败: {result.message}")
            return None, False

    # ================== 阶段3: 内层优化 ==================
    
    def inner_optimization(self, x_outer):
        """
        阶段3: 内层优化 (局部微调与约束修复)
        
        针对外层优化得到的较优解，进行局部细调微调，确保约束满足并提升目标值
        """
        print("\n🔧 阶段3: 内层优化 (局部微调与约束修复)")
        
        if x_outer is None:
            return None
        
        # 功率约束检验
        xT, yT, w, h, hm = x_outer[:5]
        r_rings = sorted(x_outer[5:])
        
        positions = self.generate_layout(xT, yT, w, h, hm, r_rings)
        feasible, power, efficiency = self._check_constraints(xT, yT, w, h, hm, positions)
        
        print(f"   初始功率: {power:.2f} MW (要求 ≥{self.min_power} MW)")
        
        if power < self.min_power:
            print("🔄 功率不足，进行增外环策略...")
            # 增外环策略: 若 Eyr < 60 MW，通过增加外环镜数 (缩小 Δθk) 或增设环来提升功率
            x_outer = self._power_enhancement(x_outer)
            
        # 镜位微调
        print("🔄 镜位微调...")
        x_optimized = self._position_refinement(x_outer)
        
        # 间距修复
        print("🔄 间距修复...")
        x_final = self._spacing_repair(x_optimized)
        
        return x_final

    def _power_enhancement(self, x):
        """功率约束修复: 增加外环镜数或新增环"""
        xT, yT, w, h, hm = x[:5]
        r_rings = list(sorted(x[5:]))
        
        # 策略1: 缩小最外环角间距 (增加镜数)
        max_attempts = 5
        for attempt in range(max_attempts):
            # 在最外环附近增加新环
            r_new = min(r_rings[-1] + h + 10, 350 - h/2)
            if r_new <= 350 - h/2:
                r_rings.append(r_new)
            
            # 重新计算功率
            positions = self.generate_layout(xT, yT, w, h, hm, r_rings)
            feasible, power, efficiency = self._check_constraints(xT, yT, w, h, hm, positions)
            
            print(f"     尝试 {attempt+1}: 功率 = {power:.2f} MW")
            
            if power >= self.min_power:
                break
        
        # 更新解向量
        x_new = list(x[:5])
        x_new.extend(r_rings[:len(x)-5])  # 保持原维度
        
        return np.array(x_new)

    def _position_refinement(self, x):
        """镜位微调: 在 (xT, yT) 附近5m范围内调整位置，选择最大的ε_yr"""
        xT, yT, w, h, hm = x[:5]
        
        best_x = x.copy()
        best_efficiency = 0
        
        # 搜索范围
        search_range = 5.0  # 5m
        search_steps = 5
        
        for dx in np.linspace(-search_range, search_range, search_steps):
            for dy in np.linspace(-search_range, search_range, search_steps):
                xT_new = xT + dx
                yT_new = yT + dy
                
                # 边界检查
                if xT_new**2 + yT_new**2 > 350**2:
                    continue
                
                # 计算效率
                r_rings = sorted(x[5:])
                positions = self.generate_layout(xT_new, yT_new, w, h, hm, r_rings)
                feasible, power, efficiency = self._check_constraints(xT_new, yT_new, w, h, hm, positions)
                
                if feasible and efficiency > best_efficiency:
                    best_efficiency = efficiency
                    best_x[0] = xT_new
                    best_x[1] = yT_new
        
        return best_x

    def _spacing_repair(self, x):
        """间距修复: 对违反间距约束rk+1-rk ≥ h+5的环，调整rk或移除rk使间距约束满足"""
        xT, yT, w, h, hm = x[:5]
        r_rings = sorted(x[5:])
        
        # 检查并修复间距
        fixed_rings = [r_rings[0]]  # 保持第一环
        
        for i in range(1, len(r_rings)):
            min_spacing = h + 5
            min_radius = fixed_rings[-1] + min_spacing
            
            if r_rings[i] >= min_radius:
                fixed_rings.append(r_rings[i])
            else:
                # 调整半径或跳过
                adjusted_radius = min_radius
                if adjusted_radius <= 350 - h/2:
                    fixed_rings.append(adjusted_radius)
        
        # 更新解向量
        x_fixed = x.copy()
        for i, r in enumerate(fixed_rings):
            if i < len(x) - 5:
                x_fixed[5 + i] = r
        
        return x_fixed

    # ================== 阶段4: 光学效率计算 ==================
    
    def calculate_final_efficiency(self, x):
        """
        阶段4: 光学效率计算 (复用问题一模块)
        
        调用问题一验证的效率计算模块，批量计算年均性能
        """
        print("\n📊 阶段4: 光学效率计算 (复用问题一模块)")
        
        xT, yT, w, h, hm = x[:5]
        r_rings = sorted(x[5:])
        
        # 生成最终布局
        positions = self.generate_layout(xT, yT, w, h, hm, r_rings)
        
        print(f"   定日镜数量: {len(positions)}")
        print(f"   镜面总面积: {len(positions) * w * h:.0f} m²")
        
        # 太阳位置与DNI
        # 余弦效率 η_cos = cos θ (θ为太阳入射角)
        # 大气透射率 η_at = 0.99321 - 0.0001176*d_HR + 1.97×10^-8*d²_HR 
        # 阴影遮挡效率 η_sb ≈ 1 - exp(-β²_f/(2σ²)) (β_f为收光方向偏角)
        
        total_power = 0.0
        total_efficiency = 0.0
        
        print("🔄 计算年平均光学效率...")
        
        # 并行计算提升速度
        def calculate_time_efficiency(args):
            alpha_s, gamma_s = args
            time_power = 0.0
            time_efficiency = 0.0
            
            for xi, yi in positions:
                eta = self.efficiency_calc.calculate_efficiency_single(
                    xi, yi, xT, yT, self.tower_height, hm, w, h, alpha_s, gamma_s
                )
                
                power_mirror = self.dni_avg * w * h * eta / 1e6  # MW
                time_power += power_mirror
                time_efficiency += eta
            
            return time_power, time_efficiency
        
        # 使用多进程加速 (16核CPU可计算60时刻并行)
        num_processes = min(cpu_count(), 16)
        
        try:
            with Pool(processes=num_processes) as pool:
                results = pool.map(calculate_time_efficiency, self.solar_positions)
            
            # 汇总结果
            for time_power, time_efficiency in results:
                total_power += time_power
                total_efficiency += time_efficiency
        
        except Exception as e:
            print(f"⚠️  多进程计算失败，使用单进程: {e}")
            # 备用单进程计算
            for alpha_s, gamma_s in self.solar_positions:
                time_power, time_efficiency = calculate_time_efficiency((alpha_s, gamma_s))
                total_power += time_power
                total_efficiency += time_efficiency
        
        # 年平均值计算
        avg_power = total_power / len(self.solar_positions)
        avg_efficiency = total_efficiency / (len(self.solar_positions) * len(positions))
        
        print(f"✅ 光学效率计算完成")
        print(f"   年平均功率: {avg_power:.2f} MW")
        print(f"   平均光学效率: {avg_efficiency:.4f}")
        
        return avg_power, avg_efficiency, positions

    # ================== 主求解流程 ==================
    
    def solve_problem2(self):
        """问题2完整求解流程"""
        print("\n" + "="*80)
        print("🏗️  问题2求解 - 同尺寸定日镜场的优化设计")
        print("="*80)
        
        start_time = time.time()
        
        try:
            # 阶段1+2: 外层优化
            x_outer, success = self.outer_optimization()
            if not success:
                print("❌ 外层优化失败")
                return None
            
            # 阶段3: 内层优化
            x_optimized = self.inner_optimization(x_outer)
            if x_optimized is None:
                print("❌ 内层优化失败")
                return None
            
            # 阶段4: 精确效率计算
            final_power, final_efficiency, final_positions = self.calculate_final_efficiency(x_optimized)
            
            # 结果验证
            xT, yT, w, h, hm = x_optimized[:5]
            
            if final_power >= self.min_power:
                solve_time = time.time() - start_time
                
                result = {
                    'solution': x_optimized,
                    'power': final_power,
                    'efficiency': final_efficiency,
                    'num_mirrors': len(final_positions),
                    'positions': final_positions,
                    'solve_time': solve_time
                }
                
                print(f"\n🎉 问题2求解成功! (用时: {solve_time:.1f}s)")
                return result
            else:
                print(f"❌ 最终功率 {final_power:.2f} MW < {self.min_power} MW")
                return None
                
        except Exception as e:
            print(f"💥 求解过程出错: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def visualize_result(self, result):
        """可视化最优解"""
        if result is None:
            return
            
        xT, yT, w, h, hm = result['solution'][:5]
        positions = result['positions']
        
        plt.figure(figsize=(12, 10))
        
        # 绘制场地边界
        circle = plt.Circle((0, 0), self.R_field, fill=False, color='black', linewidth=2, label='场地边界')
        plt.gca().add_patch(circle)
        
        # 绘制禁区
        keepout = plt.Circle((0, 0), self.R_keepout, fill=False, color='red', linewidth=2, label='禁区')
        plt.gca().add_patch(keepout)
        
        # 绘制定日镜
        for xi, yi in positions:
            rect = plt.Rectangle((xi-w/2, yi-h/2), w, h, 
                               fill=True, color='blue', alpha=0.6, linewidth=0.5)
            plt.gca().add_patch(rect)
        
        # 绘制吸收塔
        plt.plot(xT, yT, 'ro', markersize=10, label=f'吸收塔 ({xT:.1f}, {yT:.1f})')
        
        plt.axis('equal')
        plt.grid(True, alpha=0.3)
        plt.legend()
        plt.title(f'问题2最优解 - 功率: {result["power"]:.2f} MW, 镜数: {result["num_mirrors"]}')
        plt.xlabel('X (m)')
        plt.ylabel('Y (m)')
        
        plt.tight_layout()
        plt.savefig('problem2_optimal_layout.png', dpi=300, bbox_inches='tight')
        plt.show()
        
        print(f"📊 布局图已保存: problem2_optimal_layout.png") 